"""
文件处理工具模块
支持PDF和DOC/DOCX文件的内容提取
"""

import io
import logging
from typing import Tuple, Optional
from pathlib import Path

# PDF处理
try:
    import pypdf
except ImportError:
    pypdf = None

# DOC/DOCX处理
try:
    from docx import Document
except ImportError:
    Document = None

# 通用文档处理
try:
    from unstructured.partition.auto import partition
except ImportError:
    partition = None

logger = logging.getLogger(__name__)


class FileProcessor:
    """文件处理器，支持多种文档格式的内容提取"""
    
    SUPPORTED_EXTENSIONS = {'.pdf', '.doc', '.docx'}
    MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
    
    @classmethod
    def is_supported_file(cls, filename: str) -> bool:
        """检查文件是否支持处理"""
        return Path(filename).suffix.lower() in cls.SUPPORTED_EXTENSIONS
    
    @classmethod
    def extract_text_from_pdf(cls, file_content: bytes) -> Tuple[bool, str]:
        """从PDF文件中提取文本内容"""
        if pypdf is None:
            return False, "PDF处理库未安装"
        
        try:
            pdf_file = io.BytesIO(file_content)
            reader = pypdf.PdfReader(pdf_file)
            
            text_content = []
            for page_num, page in enumerate(reader.pages):
                try:
                    page_text = page.extract_text()
                    if page_text.strip():
                        text_content.append(f"--- 第{page_num + 1}页 ---\n{page_text}")
                except Exception as e:
                    logger.warning(f"提取第{page_num + 1}页内容时出错: {e}")
                    continue
            
            if not text_content:
                return False, "PDF文件中未找到可提取的文本内容"
            
            full_text = "\n\n".join(text_content)
            return True, full_text
            
        except Exception as e:
            logger.error(f"PDF文件处理失败: {e}")
            return False, f"PDF文件处理失败: {str(e)}"
    
    @classmethod
    def extract_text_from_docx(cls, file_content: bytes) -> Tuple[bool, str]:
        """从DOCX文件中提取文本内容"""
        if Document is None:
            return False, "DOCX处理库未安装"
        
        try:
            docx_file = io.BytesIO(file_content)
            doc = Document(docx_file)
            
            text_content = []
            
            # 提取段落文本
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():
                    text_content.append(paragraph.text)
            
            # 提取表格文本
            for table in doc.tables:
                table_text = []
                for row in table.rows:
                    row_text = []
                    for cell in row.cells:
                        if cell.text.strip():
                            row_text.append(cell.text.strip())
                    if row_text:
                        table_text.append(" | ".join(row_text))
                
                if table_text:
                    text_content.append("--- 表格内容 ---\n" + "\n".join(table_text))
            
            if not text_content:
                return False, "DOCX文件中未找到可提取的文本内容"
            
            full_text = "\n\n".join(text_content)
            return True, full_text
            
        except Exception as e:
            logger.error(f"DOCX文件处理失败: {e}")
            return False, f"DOCX文件处理失败: {str(e)}"
    
    @classmethod
    def extract_text_from_doc(cls, file_content: bytes) -> Tuple[bool, str]:
        """从DOC文件中提取文本内容（使用unstructured库）"""
        if partition is None:
            return False, "文档处理库未安装，无法处理DOC格式"
        
        try:
            # 将文件内容写入临时文件
            import tempfile
            with tempfile.NamedTemporaryFile(suffix='.doc', delete=False) as tmp_file:
                tmp_file.write(file_content)
                tmp_file_path = tmp_file.name
            
            try:
                # 使用unstructured处理DOC文件
                elements = partition(filename=tmp_file_path)
                text_content = []
                
                for element in elements:
                    if hasattr(element, 'text') and element.text.strip():
                        text_content.append(element.text)
                
                if not text_content:
                    return False, "DOC文件中未找到可提取的文本内容"
                
                full_text = "\n\n".join(text_content)
                return True, full_text
                
            finally:
                # 清理临时文件
                import os
                try:
                    os.unlink(tmp_file_path)
                except:
                    pass
                    
        except Exception as e:
            logger.error(f"DOC文件处理失败: {e}")
            return False, f"DOC文件处理失败: {str(e)}"
    
    @classmethod
    def extract_text(cls, filename: str, file_content: bytes) -> Tuple[bool, str]:
        """根据文件类型提取文本内容"""
        
        # 检查文件大小
        if len(file_content) > cls.MAX_FILE_SIZE:
            return False, f"文件大小超过限制（最大{cls.MAX_FILE_SIZE // (1024*1024)}MB）"
        
        # 检查文件类型
        if not cls.is_supported_file(filename):
            return False, f"不支持的文件类型。支持的格式: {', '.join(cls.SUPPORTED_EXTENSIONS)}"
        
        file_ext = Path(filename).suffix.lower()
        
        try:
            if file_ext == '.pdf':
                return cls.extract_text_from_pdf(file_content)
            elif file_ext == '.docx':
                return cls.extract_text_from_docx(file_content)
            elif file_ext == '.doc':
                return cls.extract_text_from_doc(file_content)
            else:
                return False, f"不支持的文件格式: {file_ext}"
                
        except Exception as e:
            logger.error(f"文件处理过程中发生未知错误: {e}")
            return False, f"文件处理失败: {str(e)}"


def process_uploaded_file(filename: str, file_content: bytes) -> Tuple[bool, str]:
    """
    处理上传的文件，提取文本内容
    
    Args:
        filename: 文件名
        file_content: 文件二进制内容
    
    Returns:
        Tuple[bool, str]: (是否成功, 提取的文本内容或错误信息)
    """
    return FileProcessor.extract_text(filename, file_content)


# 便捷函数
def get_supported_extensions() -> set:
    """获取支持的文件扩展名"""
    return FileProcessor.SUPPORTED_EXTENSIONS.copy()


def is_supported_file(filename: str) -> bool:
    """检查文件是否支持处理"""
    return FileProcessor.is_supported_file(filename)
